import cv2
import json
import os
from src.utils import load_dataset
from mird import compute_mird
from similarity_metric import multi_modal_cosine_similarity, find_best_match_position, save_matching_image


def load_true_positions(filepath):
    """
    从JSON文件加载真实位置。

    :param filepath: JSON文件路径。
    :return: 包含真实位置的字典。
    """
    with open(filepath, 'r') as f:
        true_positions = json.load(f)
    return true_positions


def find_true_position(filename, true_positions):
    """
    找到对应文件的真实位置。

    :param filename: 文件名。
    :param true_positions: 真实位置的列表。
    :return: 文件对应的真实位置。
    """
    for entry in true_positions:
        if entry['filename'] == filename:
            return entry['Tx'], entry['Ty']
    return 0, 0  # 默认位置，如果未找到


def load_existing_matching_results(filepath):
    """
    从JSON文件加载现有的匹配结果。

    :param filepath: JSON文件路径。
    :return: 包含现有匹配结果的列表。
    """
    if os.path.exists(filepath):
        with open(filepath, 'r') as f:
            return json.load(f)
    return []


def main():
    ref_images, temp_images = load_dataset('../data/optical_to_SAR')

    assert len(ref_images) == len(temp_images), "参考图像和模板图像的数量不一致"

    # 加载真实匹配位置
    true_positions = load_true_positions('../data/optical_to_SAR/processed/crop_positions.json')

    # 加载现有匹配结果
    matching_results_filepath = '../data/optical_to_SAR/matching_results.json'
    matching_results = load_existing_matching_results(matching_results_filepath)

    # 确定下一个image_pair_index的起始值
    if matching_results:
        next_index = max(result['image_pair_index'] for result in matching_results) + 1
    else:
        next_index = 0

    for idx, (ref_img, temp_img) in enumerate(zip(ref_images, temp_images)):
        # 转换为灰度图像进行计算
        ref_img_gray = cv2.cvtColor(ref_img, cv2.COLOR_BGR2GRAY)
        temp_img_gray = cv2.cvtColor(temp_img, cv2.COLOR_BGR2GRAY)

        print(f"正在处理图像对： {idx + 1}/{len(ref_images)}")
        ref_mird = compute_mird(ref_img_gray)
        temp_mird = compute_mird(temp_img_gray)
        print(f"图像对 {idx + 1}/{len(ref_images)} 的MIRD计算完成！")

        print(f"开始进行图像对 {idx + 1}/{len(ref_images)} 的余弦相似性度量")
        similarity_map = multi_modal_cosine_similarity(ref_mird, temp_mird)
        best_match_position = find_best_match_position(similarity_map)
        print(f"最佳匹配位置：{best_match_position}")

        original_filename = os.path.basename(os.listdir('../data/optical_to_SAR/processed/reference')[idx])
        output_filename = original_filename.replace('_s1_', '_match_').replace('_s2_', '_match_')
        save_matching_image(ref_img, best_match_position, temp_img.shape[:2], output_filename)
        print(f"匹配结果已保存到：../data/optical_to_SAR/matching/{output_filename}")

        # 从true_positions中获取真实位置
        true_filename = original_filename.replace('_s1_', '_s2_').replace('_s2_', '_s1_')
        true_position = find_true_position(true_filename, true_positions)

        # 这里需要将best_match_position的顺序从(y, x)调整为(x, y)
        best_match_x, best_match_y = best_match_position[1], best_match_position[0]

        matching_results.append({
            'image_pair_index': next_index + idx,
            'matched_position': [int(best_match_x), int(best_match_y)],
            'true_position': [int(true_position[0]), int(true_position[1])]
        })

    # 保存匹配结果到JSON文件
    with open(matching_results_filepath, 'w') as f:
        json.dump(matching_results, f, indent=4)
    print("匹配结果已保存到 matching_results.json")


if __name__ == "__main__":
    main()
